import torch
import torch.nn as nn
import torch.nn.functional as F

class ColorAwareAttention(nn.Module):
    """颜色感知注意力模块 - 基于你的色准控制经验"""
    
    def __init__(self, channels, reduction_ratio=8):
        super().__init__()
        self.channels = channels
        
        # 全局颜色特征提取
        self.global_color_pool = nn.AdaptiveAvgPool2d(1)
        
        # 颜色注意力机制
        self.color_attention = nn.Sequential(
            nn.Linear(channels, channels // reduction_ratio),
            nn.ReLU(inplace=True),
            nn.Linear(channels // reduction_ratio, channels),
            nn.Sigmoid()
        )
        
        # 颜色一致性约束
        self.color_consistency = nn.Conv2d(channels, channels, 1)
    
    def forward(self, x):
        batch_size, channels, height, width = x.size()
        
        # 提取全局颜色特征
        color_features = self.global_color_pool(x).view(batch_size, channels)
        
        # 计算颜色注意力权重
        color_weights = self.color_attention(color_features).view(batch_size, channels, 1, 1)
        
        # 应用颜色注意力
        color_enhanced = x * color_weights
        
        # 颜色一致性约束
        consistency = self.color_consistency(color_enhanced)
        
        return consistency + x

class UniformityEnhancementBlock(nn.Module):
    """均匀性增强模块 - 基于你的均匀性校正经验"""
    
    def __init__(self, channels):
        super().__init__()
        
        # 局部亮度分析
        self.local_brightness_analysis = nn.Sequential(
            nn.Conv2d(channels, channels // 4, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(channels // 4, 1, 3, padding=1),
            nn.Sigmoid()
        )
        
        # 均匀性补偿
        self.uniformity_compensation = nn.Sequential(
            nn.Conv2d(channels + 1, channels, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(channels, channels, 3, padding=1)
        )
    
    def forward(self, x):
        # 分析局部亮度变化
        brightness_map = self.local_brightness_analysis(x)
        
        # 拼接亮度信息
        brightness_aware = torch.cat([x, brightness_map], dim=1)
        
        # 均匀性补偿
        compensated = self.uniformity_compensation(brightness_aware)
        
        return compensated + x

class LowGrayOptimizationBlock(nn.Module):
    """低灰优化模块 - 基于你的低灰校正经验"""
    
    def __init__(self, channels):
        super().__init__()
        
        # 低灰特征提取
        self.low_gray_detector = nn.Sequential(
            nn.Conv2d(channels, channels // 2, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(channels // 2, 1, 3, padding=1),
            nn.Sigmoid()
        )
        
        # 低灰增强
        self.low_gray_enhancer = nn.Sequential(
            nn.Conv2d(channels + 1, channels, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(channels, channels, 3, padding=1)
        )
        
        # Gamma曲线调整（模拟你的Gamma校正经验）
        self.gamma_adjust = nn.Parameter(torch.ones(1))
    
    def forward(self, x):
        # 检测低灰区域
        low_gray_mask = self.low_gray_detector(x)
        
        # 拼接低灰信息
        low_gray_aware = torch.cat([x, low_gray_mask], dim=1)
        
        # 低灰增强
        enhanced = self.low_gray_enhancer(low_gray_aware)
        
        # Gamma调整（针对低灰优化）
        gamma_corrected = torch.pow(enhanced + 1e-8, self.gamma_adjust)
        
        return gamma_corrected + x

class DisplayAwareRRDBNet(nn.Module):
    """显示感知的Real-ESRGAN网络 - 基于你的全链路经验"""
    
    def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_block=16, num_grow_ch=32):
        super().__init__()
        
        self.num_feat = num_feat
        self.num_block = num_block
        
        # 初始特征提取
        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
        
        # 显示感知残差块
        self.body = nn.ModuleList()
        for i in range(num_block):
            self.body.append(
                DisplayAwareResidualBlock(num_feat + i * num_grow_ch, num_grow_ch)
            )
        
        # 特征融合
        self.conv_body = nn.Conv2d(num_feat + num_block * num_grow_ch, num_feat, 3, 1, 1)
        
        # 上采样
        self.upconv1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.upconv2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
        
        # 激活函数
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
    
    def forward(self, x):
        # 初始特征提取
        feat = self.lrelu(self.conv_first(x))
        
        # 显示感知残差块
        for block in self.body:
            feat = block(feat)
        
        # 特征融合
        feat = self.lrelu(self.conv_body(feat))
        
        # 上采样
        feat = self.lrelu(self.upconv1(F.interpolate(feat, scale_factor=2, mode='nearest')))
        feat = self.lrelu(self.upconv2(F.interpolate(feat, scale_factor=2, mode='nearest')))
        
        # 输出
        out = self.conv_last(self.lrelu(self.conv_hr(feat)))
        
        return out

class DisplayAwareResidualBlock(nn.Module):
    """显示感知残差块"""
    
    def __init__(self, in_channels, grow_channels):
        super().__init__()
        
        out_channels = in_channels + grow_channels
        
        # 基础卷积层
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, 1)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, 1)
        
        # 用于残差连接的1x1卷积，确保通道数匹配
        self.conv_identity = nn.Conv2d(in_channels, out_channels, 1, 1, 0)
        
        # 显示特性增强模块
        self.color_attention = ColorAwareAttention(out_channels)
        self.uniformity_enhance = UniformityEnhancementBlock(out_channels)
        self.low_gray_optimize = LowGrayOptimizationBlock(out_channels)
        
        # 激活函数
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
    
    def forward(self, x):
        # 提取identity并调整通道数
        identity = self.conv_identity(x)
        
        # 特征提取
        out = self.lrelu(self.conv1(x))
        out = self.conv2(out)
        
        # 显示特性增强
        out = self.color_attention(out)
        out = self.uniformity_enhance(out)
        out = self.low_gray_optimize(out)
        
        # 残差连接 - 现在通道数已匹配
        out += identity
        
        return self.lrelu(out)

# 测试网络
def test_display_aware_network():
    """测试显示感知网络"""
    model = DisplayAwareRRDBNet(num_feat=32, num_block=8)  # 轻量版测试
    
    # 测试输入
    test_input = torch.randn(1, 3, 64, 64)
    output = model(test_input)
    
    print(f"输入尺寸: {test_input.shape}")
    print(f"输出尺寸: {output.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M")
    
    return model

if __name__ == "__main__":
    model = test_display_aware_network()